#pragma once

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vector>
#include <memory>
#include <iostream>
#include <sstream>
#include <string>

// 全局节点ID计数器，用于生成唯一的节点标识
static int node_id_counter = 0;
class ExprAST;
// AST节点基类
class BaseAST {
public:
    virtual ~BaseAST() = default;
    
    virtual std::string DumpDOT() const = 0;

    virtual int GetNodeID() const = 0;
    
    static int NewNodeID() {
        return node_id_counter++;
    }
};

// 类型枚举
enum Type {
    TYPE_INT,
    TYPE_VOID
};

// 编译单元AST节点
class CompUnitAST : public BaseAST {
private:
    int node_id;

public:
    std::unique_ptr<BaseAST> decl,func_def1,func_def2;
    std::unique_ptr<BaseAST> main_func_def;

    CompUnitAST() : node_id(NewNodeID()) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        // 定义此节点
        ss << "  node" << node_id << " [label=\"CompUnit\", shape=box, style=filled, fillcolor=lightblue];\n";
        
        // 连接子节点
        if (decl) {
            ss << "  node" << node_id << " -> node" << decl->GetNodeID() 
               << " [label=\"decl\"];\n";
            ss << decl->DumpDOT();
        }
        // 连接子节点
        if (func_def1) {
            ss << "  node" << node_id << " -> node" << func_def1->GetNodeID() 
               << " [label=\"func_def\"];\n";
            ss << func_def1->DumpDOT();
        }
        

        // 连接子节点
        if (func_def2) {
            ss << "  node" << node_id << " -> node" << func_def2->GetNodeID() 
               << " [label=\"func_def\"];\n";
            ss << func_def2->DumpDOT();
        }
        if (main_func_def) {
            ss << "  node" << node_id << " -> node" << main_func_def->GetNodeID() 
               << " [label=\"main_func_def\"];\n";
            ss << main_func_def->DumpDOT();
        }
        
        return ss.str();
    }
};

// 声明AST节点
class DeclAST : public BaseAST {
private:
    int node_id;

public:
    enum DeclKind {
        CONST_DECL,
        VAR_DECL
    };

    DeclKind kind;
    std::unique_ptr<BaseAST> decl;  // 可以是 ConstDecl 或 VarDecl

    DeclAST() : node_id(NewNodeID()) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        ss << "  node" << node_id << " [label=\"Decl\", shape=box];\n";
        
        if (decl) {
            ss << "  node" << node_id << " -> node" << decl->GetNodeID() 
               << " [label=\"decl\"];\n";
            ss << decl->DumpDOT();
        }
        
        return ss.str();
    }
};

// 常量声明AST节点
class ConstDeclAST : public BaseAST {
private:
    int node_id;

public:
    Type type;
    std::string ident;
    int value;

    ConstDeclAST() : node_id(NewNodeID()) {}
    
    // 从C风格字符串构造
    ConstDeclAST(const char* id) : node_id(NewNodeID()), ident(id) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        std::string label = "ConstDecl\\nident: " + ident + 
                            "\\ntype: " + (type == TYPE_INT ? "int" : "void") +
                            "\\nvalue: " + std::to_string(value);
        
        ss << "  node" << node_id << " [label=\"" << label << "\", shape=box, style=filled, fillcolor=lightyellow];\n";
        
        return ss.str();
    }
};

// 变量声明AST节点
class VarDeclAST : public BaseAST {
private:
    int node_id;

public:
    Type type;
    std::string ident;

    VarDeclAST() : node_id(NewNodeID()) {}
    
    // 从C风格字符串构造
    VarDeclAST(const char* id) : node_id(NewNodeID()), ident(id) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        std::string label = "VarDecl\\nident: " + ident + 
                            "\\ntype: " + (type == TYPE_INT ? "int" : "void");
        
        ss << "  node" << node_id << " [label=\"" << label << "\", shape=box, style=filled, fillcolor=lightgreen];\n";
        
        return ss.str();
    }
};

// 主函数定义AST节点
class MainFuncDefAST : public BaseAST {
private:
    int node_id;

public:
    std::unique_ptr<BaseAST> block;

    MainFuncDefAST() : node_id(NewNodeID()) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        ss << "  node" << node_id << " [label=\"MainFuncDef\\nint main()\", shape=box, style=filled, fillcolor=lightcoral];\n";
        
        if (block) {
            ss << "  node" << node_id << " -> node" << block->GetNodeID() 
               << " [label=\"block\"];\n";
            ss << block->DumpDOT();
        }
        
        return ss.str();
    }
};

// 函数定义AST节点
class FuncDefAST : public BaseAST {
    private:
        int node_id;
    
    public:
        std::unique_ptr<BaseAST> block;
        Type type;
        std::string ident;
        Type param_type;
        std::string param_ident;
    
        FuncDefAST() : node_id(NewNodeID()) {}
    
        int GetNodeID() const override {
            return node_id;
        }
    
        std::string DumpDOT() const override {
            std::stringstream ss;
            
            ss << "  node" << node_id << " [label=\"FuncDefAST\\nint"
                << ident
                << "(int "
                << param_ident
                << ")\", shape=box, style=filled, fillcolor=lightcoral];\n";
            
            if (block) {
                ss << "  node" << node_id << " -> node" << block->GetNodeID() 
                   << " [label=\"block\"];\n";
                ss << block->DumpDOT();
            }
            
            return ss.str();
        }
};

// 代码块AST节点
class BlockAST : public BaseAST {
private:
    int node_id;

public:
    std::vector<std::unique_ptr<BaseAST>> items;

    BlockAST() : node_id(NewNodeID()) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        ss << "  node" << node_id << " [label=\"Block\", shape=box, style=filled, fillcolor=lightgrey];\n";
        
        // 连接所有子项目
        for (size_t i = 0; i < items.size(); i++) {
            if (items[i]) {
                ss << "  node" << node_id << " -> node" << items[i]->GetNodeID() 
                   << " [label=\"item" << i << "\"];\n";
                ss << items[i]->DumpDOT();
            }
        }
        
        return ss.str();
    }
};

// 块项目AST节点
class BlockItemAST : public BaseAST {
private:
    int node_id;

public:
    std::unique_ptr<BaseAST> item;  // 可以是 Decl 或 Stmt

    BlockItemAST() : node_id(NewNodeID()) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        ss << "  node" << node_id << " [label=\"BlockItem\", shape=box];\n";
        
        if (item) {
            ss << "  node" << node_id << " -> node" << item->GetNodeID() 
               << " [label=\"item\"];\n";
            ss << item->DumpDOT();
        }
        
        return ss.str();
    }
};

// 表达式基类
class ExprAST : public BaseAST {
    protected:
        int node_id;
    
    public:
        ExprAST() : node_id(NewNodeID()) {}
    
        int GetNodeID() const override {
            return node_id;
        }
    
        virtual std::string GetLabel() const {
            return "Expr";
        }
    
        std::string DumpDOT() const override {
            std::stringstream ss;
            ss << "  node" << node_id << " [label=\"" << GetLabel() << "\", shape=ellipse, style=filled, fillcolor=lightpink];\n";
            return ss.str();
        }
    };

// 语句AST节点
class StmtAST : public BaseAST {
private:
    int node_id;

public:
    enum StmtKind {
        ASSIGNMENT,
        BLOCK,
        RETURN,
        EMPTY,
        IF,
        PRINTF
    };

    StmtKind kind;
    std::unique_ptr<BaseAST> content;  // Assignment, Block 或 返回值标识符
    std::string return_ident;  // 如果是RETURN类型，存储返回的标识符
    int return_num;
 // 用于IF语句的额外字段
 std::unique_ptr<ExprAST> condition;  // 条件表达式
 std::unique_ptr<StmtAST> if_branch;  // if分支
 std::unique_ptr<StmtAST> else_branch;  // else分支
    //FOR printf
    std::string str="",ident="";

    StmtAST() : node_id(NewNodeID()) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        std::string label;
        
        switch (kind) {
            case ASSIGNMENT:
                label = "Stmt\\nkind: ASSIGNMENT";
                break;
            case BLOCK:
                label = "Stmt\\nkind: BLOCK";
                break;
            case RETURN:
                label = "Stmt\\nkind: RETURN\\nvalue: " + std::to_string(return_num);
                break;
            case EMPTY:
                label = "Stmt\\nkind: EMPTY";
                break;
            case IF:
                label = "Stmt\\nkind: IF";
                break;
            case PRINTF:
                  // 使用 Lambda 函数进行 % 的转义
                auto escapePercent = [](const std::string& input) -> std::string {
                    std::string result;
                    for (char c : input) {
                        if (c == '%') {
                            result += "%";
                        } else if('\"' ==c ){
                            result += "\\\"";
                        }else {
                            result += c;
                        }
                    }
                    return result;
                };
                label = "Stmt\\nkind: PRINTF\\nparam: " + escapePercent(str) + (""==ident?"":","+ident);
                break;
        }
        
        ss << "  node" << node_id << " [label=\"" << label << "\", shape=box, style=filled, fillcolor=lightskyblue];\n";
        
         if (kind == IF) {
            if (condition) {
                ss << "  node" << node_id << " -> node" << condition->GetNodeID() << " [label=\"condition\"];\n";
                ss << condition->DumpDOT();
            }
            if (if_branch) {
                ss << "  node" << node_id << " -> node" << if_branch->GetNodeID() << " [label=\"if_branch\"];\n";
                ss << if_branch->DumpDOT();
            }
            if (else_branch) {
                ss << "  node" << node_id << " -> node" << else_branch->GetNodeID() << " [label=\"else_branch\"];\n";
                ss << else_branch->DumpDOT();
            }
        }else if (content) {
            std::string edge_label;
            switch (kind) {
                case ASSIGNMENT:
                    edge_label = "assignment";
                    break;
                case BLOCK:
                    edge_label = "block";
                    break;
                default:
                    edge_label = "content";
                    break;
            }
            
            ss << "  node" << node_id << " -> node" << content->GetNodeID() 
               << " [label=\"" << edge_label << "\"];\n";
            ss << content->DumpDOT();
        }
        
        return ss.str();
    }
};

// 赋值语句AST节点
class AssignmentAST : public BaseAST {
private:
    int node_id;

public:
    std::string ident,ident_r="";
    bool is_getint=false;
    int num;

    AssignmentAST() : node_id(NewNodeID()) {}
    
    // 从C风格字符串构造
    AssignmentAST(const char* id) : node_id(NewNodeID()), ident(id) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        std::string label = "Assignment\\n" + ident + "=" 
                            + (is_getint?"getint()":(ident_r==""?std::to_string(num):ident_r)) 
                            + "\n";
        
        ss << "  node" << node_id << " [label=\"" << label << "\", shape=box, style=filled, fillcolor=lightsalmon];\n";
        
        return ss.str();
    }
};

// 数字AST节点
class NumberAST : public BaseAST {
private:
    int node_id;

public:
    int value;

    NumberAST() : node_id(NewNodeID()) {}

    NumberAST(int val) : node_id(NewNodeID()), value(val) {}

    int GetNodeID() const override {
        return node_id;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        
        std::string label = "Number\\nvalue: " + std::to_string(value);
        
        ss << "  node" << node_id << " [label=\"" << label << "\", shape=ellipse, style=filled, fillcolor=lightpink];\n";
        
        return ss.str();
    }
};



class BinaryExprAST : public ExprAST {
public:
    std::unique_ptr<ExprAST> left;
    std::unique_ptr<ExprAST> right;
    std::string op;

    BinaryExprAST() {}

    std::string GetLabel() const override {
        return "BinaryExpr\\nop: " + op;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        ss << ExprAST::DumpDOT();
        if (left) {
            ss << "  node" << node_id << " -> node" << left->GetNodeID() << " [label=\"left\"];\n";
            ss << left->DumpDOT();
        }
        if (right) {
            ss << "  node" << node_id << " -> node" << right->GetNodeID() << " [label=\"right\"];\n";
            ss << right->DumpDOT();
        }
        return ss.str();
    }
};

class UnaryExprAST : public ExprAST {
public:
    std::unique_ptr<ExprAST> operand;
    std::string op;

    UnaryExprAST() {}

    std::string GetLabel() const override {
        return "UnaryExpr\\nop: " + op;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        ss << ExprAST::DumpDOT();
        if (operand) {
            ss << "  node" << node_id << " -> node" << operand->GetNodeID() << " [label=\"operand\"];\n";
            ss << operand->DumpDOT();
        }
        return ss.str();
    }
};

class FunctionCallExprAST : public ExprAST {
public:
    std::string ident;
    std::vector<std::unique_ptr<ExprAST>> args;

    FunctionCallExprAST(const std::string& i) : ident(i) {}

    std::string GetLabel() const override {
        return "FunctionCallExprAST\\nident: " + ident;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        ss << ExprAST::DumpDOT();
        for (size_t i = 0; i < args.size(); ++i) {
            ss << "  node" << node_id << " -> node" << args[i]->GetNodeID() << " [label=\"arg" << i << "\"];\n";
            ss << args[i]->DumpDOT();
        }
        return ss.str();
    }
};

// 标识符表达式
class IdentExprAST : public ExprAST {
public:
    std::string ident;

    IdentExprAST(const std::string& id) : ident(id) {}

    std::string GetLabel() const override {
        return "IdentExpr\\nident: " + ident;
    }

    std::string DumpDOT() const override {
        std::stringstream ss;
        ss << ExprAST::DumpDOT();
        return ss.str();
    }
};